import torch
import json
from torchsummary import summary
from models.price_predictor import PricePredictor
from datasets.dataset import TicketDataset  # 用于自动获取输入维度

def print_model_structure():
    """加载模型并打印其结构，参数在函数内直接配置"""
    # ====================== 配置参数 - 可根据需要修改 ======================
    CONFIG_PATH = "movie-ticket-bidding/config/FeedForwardModel.json"
    DATA_PATH = "movie-ticket-bidding/data/processed/test.pkl"  # 用于获取输入维度
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    # ====================================================================
    
    # 检查设备
    print(f"使用设备: {DEVICE}")
    if DEVICE == "cuda":
        print(f"CUDA设备: {torch.cuda.get_device_name(0)}")
    
    # 从数据集中自动获取输入维度（和原代码保持一致）
    dataset = TicketDataset(DATA_PATH)
    input_dim = dataset[0][0].shape[0]
    print(f"自动获取输入特征维度: {input_dim}\n")
    
    # 加载模型配置
    with open(CONFIG_PATH, "r") as f:
        model_cfg = json.load(f)["model"]
        print(f"已加载模型配置: {CONFIG_PATH}")
        print(f"模型配置详情: {json.dumps(model_cfg, indent=2)}\n")
    
    # 创建模型实例
    model = PricePredictor(CONFIG_PATH, input_dim).to(DEVICE)
    model.eval()  # 设置为评估模式
    
    # 打印模型整体结构
    print("="*80)
    print("模型整体结构:")
    print(model)
    print("="*80 + "\n")
    
    # 使用torchsummary打印详细层信息（包括输出形状和参数数量）
    print("模型详细层信息 (输入形状: [batch_size, input_dim]):")
    # 替换 summary 部分
    print("模型详细层信息:")
    with torch.no_grad():
        # 创建测试输入 (batch_size=2, input_dim=7)
        test_input = torch.randn(2, input_dim).to(DEVICE)
        print(f"输入形状: {test_input.shape}")
        
        # 跟踪每一层的输出
        x = test_input
        for i, layer in enumerate(model.layers):
            x = layer(x)
            print(f"层 {i+1} ({layer.__class__.__name__}) 输出形状: {x.shape}")
    # summary(model, input_size=(1,input_dim), device=DEVICE)
    summary(model, input_size=(input_dim,), device=DEVICE)
    
    
    return model

if __name__ == "__main__":
    # 直接调用函数，无需命令行参数
    print_model_structure()
